{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 12.3 实现Seq2Seq中英翻译"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "L = tf.keras.layers\n",
    "\n",
    "model = tf.keras.Sequential([\n",
    "    L.Dense(12, input_dim=7, activation='relu', name='dense_layer0'),\n",
    "    L.Dense(6, activation='relu', name='dense_layer1'),\n",
    "    L.Dense(1, activation='sigmoid', name='output_layer')\n",
    "])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 定义层\n",
    "# Input 层是个特殊的层，用于定义输入的形状，初始化后是一个张量\n",
    "input_layer = L.Input(shape=(7,), name='input_layer')\n",
    "# 除了 Input 外的层对象接受张量为参数，返回一个张量\n",
    "dense_layer0 = L.Dense(12, activation='relu', name='dense_layer0')\n",
    "dense_layer1 = L.Dense(6, activation='relu', name='dense_layer1')\n",
    "output_layer = L.Dense(1, activation='sigmoid', name='output_layer')\n",
    "\n",
    "# 定义运算流程\n",
    "# input_layer 本身就是输入张量\n",
    "inputs = input_layer\n",
    "# 输入张量作为参数传入 dense_layer0 层，经过前向传播后得到张量 x\n",
    "x = dense_layer0(inputs)\n",
    "# 张量 x 作为参数传入 dense_layer1 层，经过前向传播后得到新张量 x\n",
    "x = dense_layer1(x)\n",
    "# 张量 x 作为参数传入 output_layer 层，经过前向传播后得到模型输出张量\n",
    "predictions = output_layer(x)\n",
    "\n",
    "# 使用输入和输出张量初始化一个函数式模型\n",
    "model = tf.keras.Model(inputs=inputs, outputs=predictions)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>en</th>\n",
       "      <th>cn</th>\n",
       "      <th>cc</th>\n",
       "      <th>en_cutted</th>\n",
       "      <th>cn_cutted</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Hi.</td>\n",
       "      <td>嗨。</td>\n",
       "      <td>CC-BY 2.0 (France) Attribution: tatoeba.org #5...</td>\n",
       "      <td>[hi, .]</td>\n",
       "      <td>[&lt;BOS&gt;, 嗨, 。, &lt;EOS&gt;]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>Hi.</td>\n",
       "      <td>你好。</td>\n",
       "      <td>CC-BY 2.0 (France) Attribution: tatoeba.org #5...</td>\n",
       "      <td>[hi, .]</td>\n",
       "      <td>[&lt;BOS&gt;, 你, 好, 。, &lt;EOS&gt;]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>Run.</td>\n",
       "      <td>你用跑的。</td>\n",
       "      <td>CC-BY 2.0 (France) Attribution: tatoeba.org #4...</td>\n",
       "      <td>[run, .]</td>\n",
       "      <td>[&lt;BOS&gt;, 你, 用, 跑, 的, 。, &lt;EOS&gt;]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>Wait!</td>\n",
       "      <td>等等！</td>\n",
       "      <td>CC-BY 2.0 (France) Attribution: tatoeba.org #1...</td>\n",
       "      <td>[wait, !]</td>\n",
       "      <td>[&lt;BOS&gt;, 等, 等, ！, &lt;EOS&gt;]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>Wait!</td>\n",
       "      <td>等一下！</td>\n",
       "      <td>CC-BY 2.0 (France) Attribution: tatoeba.org #1...</td>\n",
       "      <td>[wait, !]</td>\n",
       "      <td>[&lt;BOS&gt;, 等, 一, 下, ！, &lt;EOS&gt;]</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "      en     cn                                                 cc  en_cutted  \\\n",
       "0    Hi.     嗨。  CC-BY 2.0 (France) Attribution: tatoeba.org #5...    [hi, .]   \n",
       "1    Hi.    你好。  CC-BY 2.0 (France) Attribution: tatoeba.org #5...    [hi, .]   \n",
       "2   Run.  你用跑的。  CC-BY 2.0 (France) Attribution: tatoeba.org #4...   [run, .]   \n",
       "3  Wait!    等等！  CC-BY 2.0 (France) Attribution: tatoeba.org #1...  [wait, !]   \n",
       "4  Wait!   等一下！  CC-BY 2.0 (France) Attribution: tatoeba.org #1...  [wait, !]   \n",
       "\n",
       "                       cn_cutted  \n",
       "0           [<BOS>, 嗨, 。, <EOS>]  \n",
       "1        [<BOS>, 你, 好, 。, <EOS>]  \n",
       "2  [<BOS>, 你, 用, 跑, 的, 。, <EOS>]  \n",
       "3        [<BOS>, 等, 等, ！, <EOS>]  \n",
       "4     [<BOS>, 等, 一, 下, ！, <EOS>]  "
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 引入该章节全部依赖\n",
    "import collections\n",
    "import operator\n",
    "import random\n",
    "from typing import List, Dict\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from hanziconv import HanziConv\n",
    "from segtok.tokenizer import word_tokenizer\n",
    "from tensorflow import keras\n",
    "from tensorflow.keras.preprocessing.sequence import pad_sequences\n",
    "from tensorflow.keras.utils import to_categorical\n",
    "\n",
    "data_path = 'data/cmn-eng/cmn.txt'\n",
    "df = pd.read_csv(data_path, header=None, sep='\\t')\n",
    "\n",
    "# 原始数据没有表头，我们用这个方法增加表头\n",
    "df.columns = ['en', 'cn', 'cc']\n",
    "# 为了加速训练过程，我们只取前 5000 条数据，你也可以使用全部数据进行训练\n",
    "df = df[:5000]\n",
    "\n",
    "# 把繁体中文转换文简体中文\n",
    "df['cn'] = df['cn'].apply(lambda x: HanziConv.toSimplified(x))\n",
    "\n",
    "# 使用 segtok 分词，分词前把全部文本转为小写\n",
    "df['en_cutted'] = df['en'].apply(lambda x: word_tokenizer(x.lower()))\n",
    "# 基于字的分词，同时增加开始和结束标志\n",
    "df['cn_cutted'] = df['cn'].apply(lambda x: ['<BOS>'] + list(x) + ['<EOS>'])\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Processor(object):\n",
    "\n",
    "    def build_token_dict(self, corpus: List[List[str]]):\n",
    "        \"\"\"\n",
    "        构建 token 字典，这个方法将会遍历分词后的语料，构建一个标记频率字典和标记与索引的映射字典\n",
    "\n",
    "        Args:\n",
    "            corpus: 所有分词后的语料\n",
    "        \"\"\"\n",
    "        token2idx = {\n",
    "            '<PAD>': 0,\n",
    "            '<UNK>': 1,\n",
    "            '<BOS>': 2,\n",
    "            '<EOS>': 3\n",
    "        }\n",
    "\n",
    "        token2count = {}\n",
    "        for sentence in corpus:\n",
    "            for token in sentence:\n",
    "                count = token2count.get(token, 0)\n",
    "                token2count[token] = count + 1\n",
    "        # 按照词频降序排序\n",
    "        sorted_token2count = sorted(token2count.items(),\n",
    "                                    key=operator.itemgetter(1),\n",
    "                                    reverse=True)\n",
    "        token2count = collections.OrderedDict(sorted_token2count)\n",
    "\n",
    "        for token in token2count.keys():\n",
    "            if token not in token2idx:\n",
    "                token2idx[token] = len(token2idx)\n",
    "        return token2idx, token2count\n",
    "\n",
    "    @staticmethod\n",
    "    def numerize_sequences(sequence: List[str],\n",
    "                           token2index: Dict[str, int]) -> List[int]:\n",
    "        \"\"\"\n",
    "        将分词后的标记（token）数组转换成对应的索引数组\n",
    "        如 ['我', '想', '睡觉'] -> [10, 313, 233]\n",
    "\n",
    "        Args:\n",
    "            sequence: 分词后的标记数组\n",
    "            token2index: 索引词典\n",
    "        Returns: 输入数据对应的索引数组\n",
    "        \"\"\"\n",
    "        token_result = []\n",
    "        for token in sequence:\n",
    "            token_index = token2index.get(token)\n",
    "            if token_index is None:\n",
    "                token_index = token2index['<UNK>']\n",
    "            token_result.append(token_index)\n",
    "        return token_result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "p = Processor()\n",
    "\n",
    "p.input2idx, p.input2count = p.build_token_dict(df.en_cutted.to_list())\n",
    "p.output2idx, p.output2count = p.build_token_dict(df.cn_cutted.to_list())\n",
    "\n",
    "p.idx2output = dict([(v, k) for k, v in p.output2idx.items()])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "ENCODER_DIM = len(p.input2idx)\n",
    "DECODER_DIM = len(p.output2idx)\n",
    "\n",
    "# 读取序列长度，用于补全数据\n",
    "EN_SEQ_LEN = max([len(seq) for seq in df.en_cutted.to_list()])\n",
    "CN_SEQ_LEN = max([len(seq) for seq in df.cn_cutted.to_list()])\n",
    "\n",
    "# 隐藏层数量\n",
    "HIDDEN_LAYER_DIM = 512"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenized_en = []\n",
    "tokenized_cn = []\n",
    "\n",
    "for input_seq in df.en_cutted.to_list():\n",
    "    tokenized_en.append(p.numerize_sequences(input_seq, p.input2idx))\n",
    "\n",
    "for output_seq in df.cn_cutted.to_list():\n",
    "    tokenized_cn.append(p.numerize_sequences(output_seq, p.output2idx))\n",
    "\n",
    "padded_en = pad_sequences(tokenized_en, EN_SEQ_LEN, padding='post', truncating='post')\n",
    "padded_cn = pad_sequences(tokenized_cn, CN_SEQ_LEN, padding='post', truncating='post')\n",
    "\n",
    "encoder_input_data = padded_en\n",
    "# 第 0 个时长到倒数第 2 个时长的序列作为解码器输入\n",
    "decoder_input_data = padded_cn[:, :-1]\n",
    "# 第 1 个时长到最后一个时长的序列作为解码器输入\n",
    "# 由于输出层通过交叉熵计算损失，还需要把解码器输出转换为 one-hot 编码\n",
    "decoder_output_data = to_categorical(padded_cn[:, 1:], DECODER_DIM)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "L = keras.layers\n",
    "\n",
    "# 编码器输入\n",
    "encoder_inputs = L.Input(shape=(None,),\n",
    "                         name='encoder_inputs')\n",
    "\n",
    "# 编码器 Embedding 层\n",
    "encoder_embedding_layer = L.Embedding(input_dim=ENCODER_DIM,\n",
    "                                      output_dim=64,\n",
    "                                      name='encoder_embedding')\n",
    "\n",
    "# 编码器 LSTM 层\n",
    "encoder_lstm_layer = L.LSTM(HIDDEN_LAYER_DIM,\n",
    "                            return_state=True,  # 返回编码器的隐藏层状态\n",
    "                            name='encoder_lstm')\n",
    "\n",
    "encoder_embeddings = encoder_embedding_layer(encoder_inputs)\n",
    "# 获取编码器 LSTM 层的隐藏层状态\n",
    "encoder_outputs, state_h, state_c = encoder_lstm_layer(encoder_embeddings)\n",
    "\n",
    "encoder_states = [state_h, state_c]\n",
    "\n",
    "# 解码器输入\n",
    "decoder_inputs = L.Input(shape=(None,),\n",
    "                         name='decoder_inputs')\n",
    "# 解码器 Embedding 层\n",
    "decoder_embedding_layer = L.Embedding(input_dim=DECODER_DIM,\n",
    "                                      output_dim=64,\n",
    "                                      name='decoder_embedding')\n",
    "# 解码器 LSTM 层\n",
    "decoder_lstm_layer = L.LSTM(HIDDEN_LAYER_DIM,\n",
    "                            return_sequences=True,  # 返回序列\n",
    "                            return_state=True,  # 返回编码器的隐藏层状态\n",
    "                            name='decoder_lstm')\n",
    "\n",
    "# 解码器的全连接输出层\n",
    "decoder_dense_layer = L.Dense(DECODER_DIM,\n",
    "                              activation='softmax',\n",
    "                              name='decoder_dense')\n",
    "\n",
    "decoder_embeddings = decoder_embedding_layer(decoder_inputs)\n",
    "# 使用编码器隐藏层状态作为解码器的初始状态\n",
    "decoder_lstm_output, state_h, state_c = decoder_lstm_layer(decoder_embeddings,\n",
    "                                                           initial_state=encoder_states)\n",
    "decoder_outputs = decoder_dense_layer(decoder_lstm_output)\n",
    "\n",
    "# 构造模型，输入为编码器输入和解码器输入，输出为编码器输出\n",
    "model = keras.Model([encoder_inputs, decoder_inputs], decoder_outputs)\n",
    "model.compile(optimizer='rmsprop', loss='categorical_crossentropy', metrics=['acc'])\n",
    "model.summary()\n",
    "\n",
    "# 训练 100 轮次\n",
    "model.fit([encoder_input_data, decoder_input_data],\n",
    "          decoder_output_data,\n",
    "          epochs=100,\n",
    "          batch_size=64,\n",
    "          callbacks=[])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: \"model_4\"\n",
      "_________________________________________________________________\n",
      "Layer (type)                 Output Shape              Param #   \n",
      "=================================================================\n",
      "encoder_inputs (InputLayer)  [(None, None)]            0         \n",
      "_________________________________________________________________\n",
      "encoder_embedding (Embedding (None, None, 64)          142528    \n",
      "_________________________________________________________________\n",
      "encoder_lstm (LSTM)          [(None, 512), (None, 512) 1181696   \n",
      "=================================================================\n",
      "Total params: 1,324,224\n",
      "Trainable params: 1,324,224\n",
      "Non-trainable params: 0\n",
      "_________________________________________________________________\n",
      "Model: \"model_5\"\n",
      "__________________________________________________________________________________________________\n",
      "Layer (type)                    Output Shape         Param #     Connected to                     \n",
      "==================================================================================================\n",
      "decoder_inputs (InputLayer)     [(None, None)]       0                                            \n",
      "__________________________________________________________________________________________________\n",
      "decoder_embedding (Embedding)   (None, None, 64)     101824      decoder_inputs[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "input_3 (InputLayer)            [(None, 512)]        0                                            \n",
      "__________________________________________________________________________________________________\n",
      "input_4 (InputLayer)            [(None, 512)]        0                                            \n",
      "__________________________________________________________________________________________________\n",
      "decoder_lstm (LSTM)             [(None, None, 512),  1181696     decoder_embedding[0][0]          \n",
      "                                                                 input_3[0][0]                    \n",
      "                                                                 input_4[0][0]                    \n",
      "__________________________________________________________________________________________________\n",
      "decoder_dense (Dense)           (None, None, 1591)   816183      decoder_lstm[2][0]               \n",
      "==================================================================================================\n",
      "Total params: 2,099,703\n",
      "Trainable params: 2,099,703\n",
      "Non-trainable params: 0\n",
      "__________________________________________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "# 编码模型接受编码器输入，输出编码器隐藏层状态\n",
    "encoder_model = keras.Model(encoder_inputs, encoder_states)\n",
    "encoder_model.summary()\n",
    "\n",
    "# 解码模型接受解码器 Embedding 结果和上一次的隐藏层状态作为输入\n",
    "decoder_state_input_h = L.Input(shape=(HIDDEN_LAYER_DIM,))\n",
    "decoder_state_input_c = L.Input(shape=(HIDDEN_LAYER_DIM,))\n",
    "decoder_states_inputs = [decoder_state_input_h, decoder_state_input_c]\n",
    "\n",
    "decoder_lstm_outputs, h, c = decoder_lstm_layer(decoder_embeddings,\n",
    "                                                initial_state=decoder_states_inputs)\n",
    "\n",
    "# 解码器以目标序列和当前隐藏层状态作为输出\n",
    "decoder_states = [h, c]\n",
    "decoder_outputs = decoder_dense_layer(decoder_lstm_outputs)\n",
    "\n",
    "decoder_model = keras.Model([decoder_inputs] + decoder_states_inputs,\n",
    "                            [decoder_outputs] + decoder_states)\n",
    "decoder_model.summary()\n",
    "\n",
    "\n",
    "def translate_sentence(sentence: List[str]):\n",
    "    \"\"\"\n",
    "    翻译句子\n",
    "    Args:\n",
    "        sentence: 原始句子\n",
    "\n",
    "    Returns:\n",
    "        翻译结果\n",
    "    \"\"\"\n",
    "    # 输入句子转换为 idx 序列，补全序列\n",
    "    vec_sen = p.numerize_sequences(sentence, p.input2idx)\n",
    "    vec_sen = pad_sequences([vec_sen], EN_SEQ_LEN, padding='post', truncating='post')\n",
    "    # 获取 Thought Vector\n",
    "    h1, c1 = encoder_model.predict(vec_sen)\n",
    "\n",
    "    # 以开始标记 <BOS> 作为输入标记，开始预测\n",
    "    target_seq = np.array([[p.output2idx['<BOS>']]])\n",
    "\n",
    "    outputs: List[int] = []\n",
    "\n",
    "    while True:\n",
    "        # 预测下一个标记，更新隐藏层状态\n",
    "        output_tokens, h1, c1 = decoder_model.predict([target_seq, h1, c1])\n",
    "        # 通过 argmax 方法，得到下一个标记的 id\n",
    "        sampled_token_index: int = np.argmax(output_tokens[0, -1, :])\n",
    "\n",
    "        # 当标记为结束标志或者序列过长时候停止预测\n",
    "        if sampled_token_index == p.output2idx['<EOS>'] or len(outputs) > 30:\n",
    "            break\n",
    "\n",
    "        outputs.append(sampled_token_index)\n",
    "        # 使用预测标记作为下一次的输入\n",
    "        target_seq = np.array([[sampled_token_index]])\n",
    "\n",
    "    return ''.join([p.idx2output[output] for output in outputs])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "feel free to stay .           -> 欢迎留下来。\n",
      "i use twitter .               -> 我用Twitter。\n",
      "their eyes met .              -> 他们目光相接。\n",
      "please eat some cake .        -> 请吃点蛋糕。\n",
      "it can't be true .            -> 那不可能是真的。\n",
      "may we swim here ?            -> 我们能在这里游泳吗？\n",
      "yes , of course .             -> 是的，当然。\n",
      "who is that old man ?         -> 那个老男人是谁？\n",
      "you've been had .             -> 你们被骗了。\n",
      "this is a pencil .            -> 这是一支铅笔。\n"
     ]
    }
   ],
   "source": [
    "for i in range(10):\n",
    "    sentence = random.choice(df.en_cutted.to_list())\n",
    "    res = translate_sentence(sentence)\n",
    "    print(f\"{' '.join(sentence):30}-> {res}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
